﻿using System;
using System.Collections.Generic;
using JetBrains.Annotations;

namespace Tools.CollectionEx
{
    /// <summary>
    /// 优先队列
    /// </summary>
    /// <typeparam name="TElement">元素类型</typeparam>
    /// <typeparam name="TPriority">优先值类型</typeparam>
    public class PriorityQueue<TElement,TPriority>
    {
        //容器
        private IList<(TElement element, TPriority priority)> m_Nodes;
        //比较器
        private Comparison<TPriority> m_Comparison;
        /// <summary>
        /// 堆大小
        /// </summary>
        public int Size { get; private set; }

        /// <summary>
        /// 容器是否为空
        /// </summary>
        public bool Empty => Size == 0;
        /// <summary>
        /// 子节点数
        /// </summary>
        private readonly int Arity = 4;
        /// <summary>
        /// 位运算参数，和Arity 关联 具体为 Arity = 1 左移/右移(效率高点) Log2Arity
        /// 如果不是2的倍数用不上
        /// </summary>
        private const int Log2Arity = 2;

        private int GetParentIndex(int index) => (index - 1) / Arity;
        private int GetFristChildIndex(int index) => index * Arity | 1;

        /// <summary>
        /// 
        /// </summary>
        /// <param name="mComparison">比较器，不能为Null</param>
        /// <param name="arity">子树数量,必须>0</param>
        /// <exception cref="ArgumentNullException"></exception>
        public PriorityQueue([NotNull] Comparison<TPriority> mComparison, int arity = 2)
        {
            m_Comparison = mComparison ?? throw new ArgumentNullException(nameof(mComparison));
            m_Nodes = new List<(TElement, TPriority)>();
            Arity = arity;
        }

        /// <summary>
        /// 入队
        /// </summary>
        public void Enqueue(TElement element, TPriority priority)
        {
            var lastIdx = Size++;
            var pair = (element, priority);
            
            m_Nodes.Add(pair);
            MoveUpUpdate(pair, lastIdx);
        }
        
        /// <summary>
        /// 出队
        /// </summary>
        /// <returns></returns>
        public TElement Dequeue()
        {
            if (Size == 0) throw new InvalidOperationException("priority queue is Empty!");

            var element = m_Nodes[0].element;
            
            RemoveRootNode();
            return element;
        }
        
        /// <summary>
        /// 返回最优节点
        /// </summary>
        /// <returns></returns>
        public TElement Peek()
        {
            if (Size == 0) throw new InvalidOperationException("priority queue is Empty!");
            return m_Nodes[0].element;
        }

        /// <summary>
        /// 移除某个元素 On
        /// </summary>
        /// <param name="element"></param>
        /// <returns></returns>
        public bool Remove(TElement element)
        {
            if (Size == 0) return false;
            int index = -1;
            for (int i = 0; i < Size; i++)
            {
                if (element.Equals(m_Nodes[i].element))
                {
                    index = i;
                    break;
                }
            }

            if (index >= 0)
            {
                int lastNodeIndex = --Size;
                var lastNode = m_Nodes[lastNodeIndex];
                MoveDownUpdate(lastNode, lastNodeIndex);  
                m_Nodes.RemoveAt(lastNodeIndex);
            }
            return true;
        }

        public bool RemoveAt(int index)
        {
            if (Size == 0 || index < 0 || index >= Size) return false;
            int lastNodeIndex = --Size;
            var lastNode = m_Nodes[lastNodeIndex];
            MoveDownUpdate(lastNode, lastNodeIndex);
            m_Nodes.RemoveAt(lastNodeIndex);
            return true;
        }


        public int IndexOf(TElement element)
        {
            int index = -1;
            if (Size == 0) return index;
            for (int i = 0; i < Size; i++)
            {
                if (element.Equals(m_Nodes[i].element)) break;
            }

            return index;
        }

        /// <summary>
        /// 清除所有
        /// </summary>
        public void Clear() => m_Nodes.Clear();
        
        /// <summary>
        /// 移除根节点 
        /// </summary>
        private void RemoveRootNode()
        {
            int lastIndex = --Size;
            if (lastIndex > 0)
            {
                var lastNode = m_Nodes[lastIndex];
                MoveDownUpdate(lastNode, 0);
            }
            m_Nodes.RemoveAt(lastIndex);
        }
        
        /// <summary>
        /// 向上更新堆方法
        /// 从底部重新构建树形
        /// </summary>
        /// <param name="node">新添加节点</param>
        /// <param name="index">新添加节点下标</param>
        private void MoveUpUpdate((TElement element, TPriority priority) node,int index)
        {
            while (index > 0)
            {
                var parentIndex = GetParentIndex(index);
                (TElement element, TPriority priority) parent = m_Nodes[parentIndex];
                //当前节点的优先级比父节点的优先级高.父节点下放
                if (m_Comparison(node.priority, parent.priority) < 0)
                {
                    m_Nodes[index] = parent;
                    index = parentIndex;
                }
                else break;
            }
            m_Nodes[index] = node;
        }
        
        /// <summary>
        /// 向下更新堆方法
        /// 具体实现为将指定元素移到最底层
        /// </summary>
        /// <param name="node"></param>
        /// <param name="index"></param>
        private void MoveDownUpdate((TElement element, TPriority priority) node,int index)
        {
            //理论容器值，剔除了需要移除的元素
            int size = Size;
            //当前处理孩子节点下标
            int childIndex;
            //已处理的child最大优先级数据
            int maxIndex;
            (TElement element, TPriority priority) maxChild;
            (TElement element, TPriority priority) nextChildNode;
            //树的子节点数量
            int childIndexUpperBound;
            while ((childIndex = GetFristChildIndex(index))  < size)
            {
                //预设第一个
                maxIndex = childIndex;
                maxChild = m_Nodes[childIndex];
                childIndexUpperBound = Math.Min(childIndex + Arity, size);
                //处理子树，拿到最大的
                while (++childIndex < childIndexUpperBound)
                {
                    nextChildNode = m_Nodes[childIndex];
                    //如果当前节点的优先级比之前的高，替换
                    if (m_Comparison(nextChildNode.priority, maxChild.priority) <= 0)
                    {
                        maxChild = nextChildNode;
                        maxIndex = childIndex;
                    }
                }
                
                //如果子节点最大都小于传入节点优先值，退出
                if (m_Comparison(node.priority, maxChild.priority) <= 0) break;
                
                //最大的补上，继续处理他的子树
                m_Nodes[index] = maxChild;
                index = maxIndex;
            }
            //将值放在废弃位置上，也就是最后一个
            m_Nodes[index] = node;
        }

        /// <summary>
        /// 提供一个默认Int类型比较方法
        /// </summary>
        /// <param name="lhs"></param>
        /// <param name="rhs"></param>
        /// <returns></returns>
        private static int DefaultComparison(int lhs, int rhs)
        {
            return lhs - rhs;
        }
    }
}